#!/usr/bin/env python3

import argparse
import os
import sys
import typing
import binascii

def auto_int(x):
  return int(x, 0)

parser = argparse.ArgumentParser(description='Generate offline download binary')

DEFAULT_DEVICE_ID   = 0x40200001
DEFAULT_BIN_ADDRESS = 0x80000000
DEFAULT_BIN_TYPE    = 0x0
DEFAULT_ERASE_OPT   = -1
DEFAULT_PROGRAM_OPT = -1
DEFAULT_RUN_OPT     = -1

DEFAULT_BYTE = b'\xff'

RUN_OPT_JUMP  = (1 << 0)
RUN_OPT_RESET = (1 << 1)
RUN_OPT_STOP  = (1 << 2)
RUN_OPT_WAIT  = (1 << 3)

OPTION_BYTE_BASE = 0x81000000
OPTION_BYTE_SIZE = 128
SRAM_BASE        = 0x20000000

class FlashBank:
  def __init__(self, device_id, address_range):
    self.device_id     = device_id
    self.address_range = address_range

FLASH_BANKS = [
    FlashBank(0x40200001, [(0x80000000, 0x81000000), (OPTION_BYTE_BASE, OPTION_BYTE_BASE + OPTION_BYTE_SIZE)])
    ]
def is_flash(device_id, address):
  for bank in FLASH_BANKS:
    if bank.device_id == device_id:
      for range in bank.address_range:
        if address >= range[0] and address < range[1]:
          return 1
  return 0

def is_ram(device_id, address):
  if is_flash(device_id, address) != 0:
    return 0
  else:
    return 1

outgrp = parser.add_mutually_exclusive_group(required=True)
outgrp.add_argument('-o', '--output-file',     help='Output file name',   type=argparse.FileType('wb'))
outgrp.add_argument('-O', '--output-append',   help='Output append file', type=argparse.FileType('ab'))
outgrp.add_argument('-I', '--info',            help='Dump header info',   action='store_true')
parser.add_argument('-i', '--input-file',      help='Input file name',    type=argparse.FileType('rb'), default=[], nargs='+')

parser.add_argument('-d', '--device-id',       help='device ID',       type=auto_int, default=DEFAULT_DEVICE_ID)
parser.add_argument('-a', '--bin-address',     help='binary address',  type=auto_int, default=[DEFAULT_BIN_ADDRESS], nargs='+')
parser.add_argument('-t', '--bin-type',        help='binary type',     type=auto_int, default=[DEFAULT_BIN_TYPE],    nargs='+')
parser.add_argument('-e', '--erase-options',   help='erase options',   type=auto_int, default=[DEFAULT_ERASE_OPT],   nargs='+')
parser.add_argument('-p', '--program-options', help='program options', type=auto_int, default=[DEFAULT_PROGRAM_OPT], nargs='+')
parser.add_argument('-r', '--run-options',     help='run options',     type=auto_int, default=[DEFAULT_RUN_OPT],     nargs='+')

# Options to generate option bytes
parser.add_argument('--user',           help='User field in option byte',        type=auto_int)
parser.add_argument('--data0',          help='Data0 field in option byte',       type=auto_int)
parser.add_argument('--data1',          help='Data1 field in option byte',       type=auto_int)
parser.add_argument('--misc',           help='Miscellaneous configurations',     type=auto_int)
parser.add_argument('--logic-address',  help='Logic configuration address',      type=auto_int)
parser.add_argument('--logic-compress', help='Logic configuration compression',  action='store_true')
parser.add_argument('--logic-config',   help='Logic configuration file',         type=argparse.FileType('rb'))
parser.add_argument('--xmem',           help='Executable memory in option byte', type=auto_int, nargs='+')
parser.add_argument('--protect',        help='Protect flash sectors',            type=auto_int)
parser.add_argument('--lock',           help='Lock flash at the end',            action='store_true')

def to_bytes(value, length):
  return value.to_bytes(length, 'little');

def from_bytes(bytes):
  return int.from_bytes(bytes, 'little')

class DownloadHeader:
  HEADER_SIZE = 256
  def __init__(self, device_id, bin_size, bin_address, bin_type, erase_options, program_options, run_options, checksum = None):
      self.device_id       = device_id       
      self.bin_size        = bin_size        
      self.bin_address     = bin_address     
      self.bin_type        = bin_type        
      self.erase_options   = erase_options   
      self.program_options = program_options 
      self.run_options     = run_options     
      self.checksum        = checksum

  def from_bytes(header: list[bytes]):
    return DownloadHeader(
        device_id       = from_bytes(header[0:4]),
        bin_size        = from_bytes(header[4:8]),
        bin_address     = from_bytes(header[8:12]),
        bin_type        = from_bytes(header[12:16]),
        erase_options   = from_bytes(header[16:20]),
        program_options = from_bytes(header[20:24]),
        run_options     = from_bytes(header[24:28]),
        checksum        = from_bytes(header[DownloadHeader.HEADER_SIZE-4:DownloadHeader.HEADER_SIZE]))

  def to_bytes(self):
    header = bytes()
    header += to_bytes(self.device_id,       4)
    header += to_bytes(self.bin_size,        4)
    header += to_bytes(self.bin_address,     4)
    header += to_bytes(self.bin_type,        4)
    header += to_bytes(self.erase_options,   4)
    header += to_bytes(self.program_options, 4)
    header += to_bytes(self.run_options,     4)
    header += bytes.fromhex('ff' * (DownloadHeader.HEADER_SIZE - 4 - len(header)))
    header += to_bytes(binascii.crc32(header), 4)
    return header

  def get_checksum(self):
    header = self.to_bytes()
    return binascii.crc32(header[:-4])

  def is_valid(self):
    return self.get_checksum() == self.checksum

def has_post_config(args):
  return args.lock or args.protect

def get_bin_address(args, idx):
  return args.bin_address[idx if idx != None and idx < len(args.bin_address) else -1]

def get_bin_type(args, idx):
  return args.bin_type[idx if idx != None and idx < len(args.bin_type) else -1]

def get_erase_options(args, idx):
  erase_options = args.erase_options[idx if idx != None and idx < len(args.erase_options) else -1]
  if erase_options < 0:
    erase_options = is_flash(args.device_id, get_bin_address(args, idx))
  return erase_options

def get_program_options(args, idx):
  program_options = args.program_options[idx if idx != None and idx < len(args.program_options) else -1]
  if program_options < 0:
    program_options = is_flash(args.device_id, get_bin_address(args, idx))
  return program_options

def get_run_options(args, idx):
  run_options = args.run_options[idx if idx != None and idx < len(args.run_options) else -1]
  if run_options < 0:
    timeout = (-run_options)
    run_options = timeout << 16
    if is_ram(args.device_id, get_bin_address(args, idx)):
      run_options |= RUN_OPT_JUMP
    if idx == len(args.input_file) - 1 and not has_post_config(args):
      run_options |= RUN_OPT_STOP
      run_options |= RUN_OPT_RESET
  return run_options

def get_run_options_str(run_options):
  opt_str = [f"0x{run_options:x}"]
  (run_options & RUN_OPT_JUMP)  and opt_str.append("jump")
  (run_options & RUN_OPT_RESET) and opt_str.append("reset")
  (run_options & RUN_OPT_STOP)  and opt_str.append("stop")
  (run_options & RUN_OPT_WAIT)  and opt_str.append("wait")
  opt_str.append(f"timeout: {run_options >> 16}")
  return ', '.join(opt_str)

def dump_header(header: DownloadHeader):
  print(f"  Device ID:       0x{header.device_id:x}")
  print(f"  Binary Size:     {header.bin_size}")
  print(f"  Binary Address:  0x{header.bin_address:x}")
  print(f"  Binary Type:     0x{header.bin_type:x}")
  print(f"  Erase Options:   0x{header.erase_options:x}")
  print(f"  Program Options: 0x{header.program_options:x}")
  print(f"  Run Options:     {get_run_options_str(header.run_options)}")
  print(f"  Header Checksum: 0x{header.checksum:x}")

LOGIC_COMPRESS_BIN = bytes((
  0x17,0x01,0x00,0x00,0x03,0x21,0x01,0x24,0x99,0x62,0x73,0xa0,0x02,0x30,0x73,0x10,
  0x30,0x00,0xef,0x00,0x60,0x00,0x00,0x00,0x41,0x11,0x22,0xc4,0x37,0x04,0x00,0x03,
  0x26,0xc2,0x04,0x4c,0x06,0xc6,0x93,0xe7,0xf4,0x00,0x1c,0xcc,0x8d,0x20,0xb7,0x07,
  0x00,0x81,0x88,0x5f,0xdc,0x5f,0x58,0x40,0x93,0xc7,0xf7,0xff,0x63,0x08,0xf5,0x02,
  0xb7,0x07,0x00,0x03,0x84,0xcf,0xb8,0x4b,0x13,0x67,0x17,0x00,0xb8,0xcb,0xb8,0x4b,
  0x79,0x9b,0xb8,0xcb,0xb8,0x5b,0x79,0x9b,0xb8,0xdb,0xb7,0x07,0x00,0x80,0x82,0x87,
  0xb2,0x40,0x22,0x44,0x92,0x44,0x01,0x45,0x41,0x01,0x82,0x80,0x11,0x8b,0x69,0xfb,
  0x3c,0x50,0x37,0x07,0x01,0x40,0x93,0xe7,0x17,0x00,0x3c,0xd0,0x1c,0x43,0x93,0xe7,
  0x07,0x02,0x1c,0xc3,0x39,0x28,0x3c,0x50,0xf9,0x9b,0x3c,0xd0,0x55,0xbf,0xb7,0x07,
  0x00,0x03,0xd8,0x47,0x71,0x9b,0xd8,0xc7,0xd8,0x47,0x13,0x77,0xb7,0xfd,0xd8,0xc7,
  0x82,0x80,0x39,0x71,0x22,0xde,0x26,0xdc,0x4e,0xd8,0x52,0xd6,0x56,0xd4,0x5a,0xd2,
  0x5e,0xd0,0x4a,0xda,0x62,0xce,0x37,0x07,0x01,0x40,0x93,0x07,0x00,0x04,0x1c,0xc3,
  0x1c,0x41,0xfd,0x72,0x7d,0x78,0x5c,0xc7,0x50,0x41,0x93,0x82,0x02,0x81,0x16,0x91,
  0x50,0xc7,0x09,0x66,0x13,0x07,0x06,0x80,0x42,0x97,0x93,0x09,0x08,0xc0,0x03,0x49,
  0x85,0x00,0x33,0x08,0x27,0x00,0x13,0x07,0x06,0x80,0xe1,0x6e,0x4e,0x97,0x81,0x47,
  0x93,0x04,0x95,0x00,0x21,0x4f,0xfd,0x5b,0x8d,0x42,0x13,0x04,0xf0,0x1f,0xa5,0x43,
  0x81,0x45,0x13,0x0a,0x00,0x10,0x93,0x0a,0x10,0x10,0x13,0x0e,0xf0,0x0f,0xb7,0x0f,
  0x01,0x40,0x93,0x8e,0x0e,0x66,0xb3,0x08,0x27,0x00,0x37,0x0b,0x01,0x40,0x03,0xc7,
  0x04,0x00,0x22,0x09,0x21,0x0f,0x33,0x69,0x27,0x01,0x63,0x46,0x7f,0x10,0x33,0x0f,
  0x7f,0x40,0x33,0x53,0xe9,0x01,0x33,0x73,0x83,0x00,0x63,0x08,0x43,0x0d,0x63,0x0a,
  0x53,0x0f,0x13,0x85,0xf2,0x0f,0x63,0x08,0x65,0x0a,0x63,0x44,0x65,0x0e,0x63,0x5b,
  0x6e,0x0c,0x09,0x67,0x13,0x07,0x07,0x80,0x4e,0x97,0x9a,0x87,0x0a,0x97,0x93,0x87,
  0x17,0xf0,0x33,0x06,0xf8,0x00,0x8a,0x07,0x03,0x46,0x06,0x80,0xc2,0x97,0x9c,0x43,
  0x23,0x00,0xc7,0x00,0x05,0x07,0xe3,0x44,0xfe,0xfe,0x93,0xf7,0xf7,0x0f,0x23,0x00,
  0xf7,0x00,0x63,0xc2,0x0b,0x02,0x13,0x96,0x22,0x00,0x33,0x0c,0x58,0x00,0x42,0x96,
  0x23,0x20,0x76,0x01,0x23,0x00,0xfc,0x80,0x85,0x02,0x63,0x46,0x85,0x00,0x06,0x04,
  0x85,0x03,0x13,0x64,0x14,0x00,0x3e,0x86,0x85,0x05,0x62,0x06,0xa1,0x82,0x13,0xf5,
  0x35,0x00,0x7d,0x17,0xd1,0x8e,0x09,0xe5,0x23,0xa6,0xdf,0x00,0x63,0x87,0xd5,0x01,
  0x63,0x69,0x17,0x03,0x03,0x46,0x07,0x00,0xc5,0xb7,0x83,0x27,0x0b,0x01,0x05,0x45,
  0x89,0x8b,0xb1,0xe3,0x85,0x62,0x93,0x82,0x02,0x7f,0x16,0x91,0x72,0x54,0xe2,0x54,
  0x52,0x59,0xc2,0x59,0x32,0x5a,0xa2,0x5a,0x12,0x5b,0x82,0x5b,0x72,0x4c,0x21,0x61,
  0x82,0x80,0x9a,0x8b,0x1d,0xbf,0x13,0x96,0x22,0x00,0x42,0x96,0x33,0x07,0x58,0x00,
  0x23,0x20,0x76,0x01,0x23,0x00,0xf7,0x80,0xa9,0xb7,0xfd,0x5b,0x8d,0x42,0x13,0x04,
  0xf0,0x1f,0xa5,0x43,0x19,0xbf,0x03,0x25,0x0b,0x01,0x13,0x75,0x05,0x07,0x33,0x35,
  0xa0,0x00,0x4d,0xbf,0x09,0x67,0x13,0x07,0x07,0x80,0x4e,0x97,0x9a,0x87,0x0a,0x97,
  0xa9,0xb7,0x05,0x45,0x45,0xb7,0x85,0x04,0xdd,0xb5,0x00,0x00,0x00,0x00,0x00,0x00,
  0x00,0x00,0x02,0x20,0xff,0xff,0xff,0xff,0x00,0x00,0x00,0x00,
  ))
LOGIC_COMPRESS_SIZE = (len(LOGIC_COMPRESS_BIN) + 8 + 255) & ~255

def get_opt_half(value):
  return (to_bytes(value, 1) + to_bytes(0xff - value, 1)) if value != None else DEFAULT_BYTE * 2

def get_opt_dword(value):
  return (to_bytes(value, 4) + to_bytes(0xffffffff - value, 4)) if value != None else DEFAULT_BYTE * 8

def gen_option_bytes(args):
  option_bytes = bytes()

  if args.user or args.data0 or args.data1 or args.misc or args.logic_address:
    header = DownloadHeader(
        device_id       = args.device_id,
        bin_size        = OPTION_BYTE_SIZE,
        bin_address     = OPTION_BYTE_BASE,
        bin_type        = get_bin_type(args, None),
        erase_options   = get_erase_options(args, None),
        program_options = get_program_options(args, None),
        run_options     = get_run_options(args, None))
    option_bytes += header.to_bytes()
    option_bytes += get_opt_half(0xa5)
    option_bytes += get_opt_half(args.user)
    option_bytes += get_opt_half(args.data0)
    option_bytes += get_opt_half(args.data1)
    option_bytes += DEFAULT_BYTE * 24
    option_bytes += get_opt_half(args.misc)
    option_bytes += DEFAULT_BYTE * 14
    option_bytes += get_opt_dword(None if args.logic_compress else args.logic_address)
    option_bytes += get_opt_dword((args.logic_address + LOGIC_COMPRESS_SIZE) if args.logic_compress else None)
    if args.xmem:
      for xmem in args.xmem:
        option_bytes += to_bytes(xmem, 4)
    if args.logic_compress:
      option_bytes += get_opt_dword(args.logic_address)
    option_bytes += DEFAULT_BYTE * (header.bin_size + DownloadHeader.HEADER_SIZE - len(option_bytes))
  return option_bytes

def gen_logic_config(args):
  logic_config = bytes()
  if args.logic_config and args.logic_address:
    header = DownloadHeader(
        device_id       = args.device_id,
        bin_size        = os.path.getsize(args.logic_config.name) + (LOGIC_COMPRESS_SIZE if args.logic_compress else 0),
        bin_address     = args.logic_address,
        bin_type        = get_bin_type(args, None),
        erase_options   = get_erase_options(args, None),
        program_options = get_program_options(args, None),
        run_options     = get_run_options(args, None))
    logic_config += header.to_bytes()
    if args.logic_compress:
      logic_config += to_bytes(SRAM_BASE | 0x2, 4)
      logic_config += to_bytes(len(LOGIC_COMPRESS_BIN), 4)
      logic_config += LOGIC_COMPRESS_BIN
      logic_config += DEFAULT_BYTE * (LOGIC_COMPRESS_SIZE - len(LOGIC_COMPRESS_BIN) - 8)
    logic_config += args.logic_config.read()
  return logic_config

def gen_post_config(args):
  post_config = bytes()
  if args.protect:
    header = DownloadHeader(
        device_id       = args.device_id,
        bin_size        = 8,
        bin_address     = OPTION_BYTE_BASE + 8,
        bin_type        = get_bin_type(args, None),
        erase_options   = 0, # Do not erase
        program_options = get_program_options(args, None),
        run_options     = (0 if args.lock else RUN_OPT_STOP | RUN_OPT_RESET))
    post_config += header.to_bytes()
    for i in range(0, 4):
      post_config += get_opt_half((args.protect >> i * 8) & 0xff)
  if args.lock:
    header = DownloadHeader(
        device_id       = args.device_id,
        bin_size        = 2,
        bin_address     = OPTION_BYTE_BASE,
        bin_type        = get_bin_type(args, None),
        erase_options   = 0, # Do not erase
        program_options = get_program_options(args, None),
        run_options     = RUN_OPT_STOP | RUN_OPT_RESET)
    post_config += header.to_bytes()
    post_config += b'\x00\x00'
  return post_config

def write_output(args):
  if args.info:
    input_file: typing.BinaryIO = args.input_file[0]
    print(f"Input download file: {input_file.name}")
    download_bin = input_file.read()
  else:
    output_file : typing.BinaryIO = args.output_file or args.output_append
    download_bin = gen_option_bytes(args)
    download_bin += gen_logic_config(args)
    for idx in range(len(args.input_file) if args.input_file else 0):
      input_file: typing.BinaryIO = args.input_file[idx]
      bin_size = os.path.getsize(input_file.name)
      print(f"Input: {input_file.name}, {bin_size} bytes")

      bin_address     = get_bin_address(args, idx)
      bin_type        = get_bin_type(args, idx)
      erase_options   = get_erase_options(args, idx)
      program_options = get_program_options(args, idx)
      run_options     = get_run_options(args, idx)

      header = DownloadHeader(args.device_id, bin_size, bin_address, bin_type, erase_options, program_options, run_options)
      download_bin += header.to_bytes()
      download_bin += input_file.read()
    download_bin += gen_post_config(args)
    output_file.write(download_bin)
    output_file.close()
    print(f"Output: {output_file.name}, {os.path.getsize(output_file.name)} bytes")

  # Dump headers
  i = 0
  while i < len(download_bin):
    header = DownloadHeader.from_bytes(download_bin[i:])
    dump_header(header);
    if (not header.is_valid()):
      print("Error in header!")
      return -1
    i = i + DownloadHeader.HEADER_SIZE + header.bin_size
    print("")

def main():
  args = parser.parse_args(sys.argv[1:])
  return write_output(args)

if __name__ == "__main__":
  sys.exit(main())
